{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tutorial 3: Training a WSI Classification Model with ABMIL\n",
    "\n",
    "This tutorial will guide you step-by-step to train an attention-based multiple instance learning model using Trident patch embeddings. Then, we will generation attention heatmaps using the pretrained model. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A- Installation and patch feature extraction using UNI\n",
    "\n",
    "#### Step 1: Download a dataset of whole-slide images\n",
    "\n",
    "You can use your own WSIs or download a publicly available dataset, e.g. from:\n",
    "\n",
    "- **CPTAC CCRCC WSIs**: Download from the [TCIA Cancer Imaging Archive](https://www.cancerimagingarchive.net/collection/cptac-ccrcc/).\n",
    "- **Store WSIs**: Save all WSIs into a local directory, e.g.,  \n",
    "  ```bash\n",
    "  ./CPTAC-CCRCC_v1/CCRCC\n",
    "  ```\n",
    "\n",
    "#### Step 2:  Run CONCH v1.5 feature extraction:\n",
    "\n",
    "Navigate to the base directory of Trident and execute the following command:\n",
    "\n",
    "```bash\n",
    "python run_batch_of_slides.py --task all \\\n",
    "  --wsi_dir ./CPTAC-CCRCC_v1/CCRCC \\\n",
    "  --job_dir ./tutorial-3 \\\n",
    "  --patch_encoder conch_v15 \\\n",
    "  --mag 20 \\\n",
    "  --patch_size 512\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### B- Download labels with data splits\n",
    "\n",
    "Here, we use Patho-Bench CPTAC-CCRCC labels for predicting BAP1 mutation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "import pandas as pd\n",
    "\n",
    "# Download labels as csv\n",
    "datasets.load_dataset(\n",
    "    'MahmoodLab/Patho-Bench', \n",
    "    cache_dir='./tutorial-3',\n",
    "    dataset_to_download='cptac_ccrcc',     \n",
    "    task_in_dataset='BAP1_mutation',           \n",
    "    trust_remote_code=True\n",
    ")\n",
    "\n",
    "# Visualize my labels and splits\n",
    "df = pd.read_csv('tutorial-3/cptac_ccrcc/BAP1_mutation/k=all.tsv', sep=\"\\t\")\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the label distribution\n",
    "df_counts = df['BAP1_mutation'].value_counts().reset_index()\n",
    "df_counts.columns = ['BAP1_mutation', 'Count']\n",
    "df_counts\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### C- Training an ABMIL model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import h5py\n",
    "import pandas as pd\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "from trident.slide_encoder_models import ABMILSlideEncoder\n",
    "\n",
    "# Set deterministic behavior\n",
    "SEED = 1234\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.cuda.manual_seed_all(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False\n",
    "\n",
    "\n",
    "class BinaryClassificationModel(nn.Module):\n",
    "    def __init__(self, input_feature_dim=768, n_heads=1, head_dim=512, dropout=0., gated=True, hidden_dim=256):\n",
    "        super().__init__()\n",
    "        self.feature_encoder = ABMILSlideEncoder(\n",
    "            freeze=False,\n",
    "            input_feature_dim=input_feature_dim, \n",
    "            n_heads=n_heads, \n",
    "            head_dim=head_dim, \n",
    "            dropout=dropout, \n",
    "            gated=gated\n",
    "        )\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(input_feature_dim, hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_dim, 1)\n",
    "        )\n",
    "\n",
    "    def forward(self, x, return_raw_attention=False):\n",
    "        if return_raw_attention:\n",
    "            features, attn = self.feature_encoder(x, return_raw_attention=True)\n",
    "        else:\n",
    "            features = self.feature_encoder(x)\n",
    "        logits = self.classifier(features).squeeze(1)\n",
    "        \n",
    "        if return_raw_attention:\n",
    "            return logits, attn\n",
    "        \n",
    "        return logits\n",
    "\n",
    "# Initialize model\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = BinaryClassificationModel().to(device)\n",
    "\n",
    "# Custom dataset\n",
    "class H5Dataset(Dataset):\n",
    "    def __init__(self, feats_path, df, split, num_features=512):\n",
    "        self.df = df[df[\"fold_0\"] == split]\n",
    "        self.feats_path = feats_path\n",
    "        self.num_features = num_features\n",
    "        self.split = split\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        row = self.df.iloc[idx]\n",
    "        with h5py.File(os.path.join(self.feats_path, row['slide_id'] + '.h5'), \"r\") as f:\n",
    "            features = torch.from_numpy(f[\"features\"][:])\n",
    "\n",
    "        if self.split == 'train':\n",
    "            num_available = features.shape[0]\n",
    "            if num_available >= self.num_features:\n",
    "                indices = torch.randperm(num_available, generator=torch.Generator().manual_seed(SEED))[:self.num_features]\n",
    "            else:\n",
    "                indices = torch.randint(num_available, (self.num_features,), generator=torch.Generator().manual_seed(SEED))  # Oversampling\n",
    "            features = features[indices]\n",
    "\n",
    "        label = torch.tensor(row[\"BAP1_mutation\"], dtype=torch.float32)\n",
    "        return features, label\n",
    "\n",
    "# Create dataloaders\n",
    "feats_path = './tutorial-3/cptac_ccrcc/20x_512px_0px_overlap/features_conch_v15'\n",
    "batch_size = 8\n",
    "train_loader = DataLoader(H5Dataset(feats_path, df, \"train\"), batch_size=batch_size, shuffle=True, worker_init_fn=lambda _: np.random.seed(SEED))\n",
    "test_loader = DataLoader(H5Dataset(feats_path, df, \"test\"), batch_size=1, shuffle=False, worker_init_fn=lambda _: np.random.seed(SEED))\n",
    "\n",
    "# Training setup\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=4e-4)\n",
    "\n",
    "# Training loop\n",
    "num_epochs = 20\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    total_loss = 0.\n",
    "    for features, labels in train_loader:\n",
    "        features, labels = {'features': features.to(device)}, labels.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(features)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "    print(f\"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/len(train_loader):.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### D- Evaluating the ABMIL model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluation\n",
    "model.eval()\n",
    "all_labels, all_outputs = [], []\n",
    "correct = 0\n",
    "total = 0\n",
    "\n",
    "with torch.no_grad():\n",
    "    for features, labels in test_loader:\n",
    "        features, labels = {'features': features.to(device)}, labels.to(device)\n",
    "        outputs = model(features)\n",
    "        \n",
    "        # Convert logits to probabilities and binary predictions\n",
    "        predicted = (outputs > 0).float()  # Since BCEWithLogitsLoss expects raw logits\n",
    "        correct += (predicted == labels).sum().item()\n",
    "        total += labels.size(0)\n",
    "\n",
    "        all_outputs.append(outputs.cpu().numpy())  \n",
    "        all_labels.append(labels.cpu().numpy())\n",
    "\n",
    "# Compute AUC\n",
    "all_outputs = np.concatenate(all_outputs)\n",
    "all_labels = np.concatenate(all_labels)\n",
    "auc = roc_auc_score(all_labels, all_outputs)\n",
    "\n",
    "# Compute accuracy\n",
    "accuracy = correct / total\n",
    "print(f\"Test AUC: {auc:.4f}\")\n",
    "print(f\"Test Accuracy: {accuracy:.4f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### E- Extract attention heatmap for the freshly trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trident import OpenSlideWSI, visualize_heatmap\n",
    "from trident.segmentation_models import segmentation_model_factory\n",
    "from trident.patch_encoder_models import encoder_factory as patch_encoder_factory\n",
    "\n",
    "# a. Load WSI to process\n",
    "job_dir = './tutorial-3/heatmap_viz'\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "slide = OpenSlideWSI(slide_path='./CPTAC-CCRCC_v1/CCRCC/C3L-00418-22.svs', lazy_init=False)\n",
    "\n",
    "# b. Run segmentation \n",
    "segmentation_model = segmentation_model_factory(\"hest\")\n",
    "geojson_contours = slide.segment_tissue(segmentation_model=segmentation_model, job_dir=job_dir, device=device)\n",
    "\n",
    "# c. Run patch coordinate extraction\n",
    "coords_path = slide.extract_tissue_coords(\n",
    "    target_mag=20,\n",
    "    patch_size=512,\n",
    "    save_coords=job_dir,\n",
    "    overlap=256, \n",
    ")\n",
    "\n",
    "# d. Run patch feature extraction\n",
    "patch_encoder = patch_encoder_factory(\"conch_v15\").eval().to(device)\n",
    "patch_features_path = slide.extract_patch_features(\n",
    "    patch_encoder=patch_encoder,\n",
    "    coords_path=coords_path,\n",
    "    save_features=os.path.join(job_dir, f\"features_conch_v15\"),\n",
    "    device=device\n",
    ")\n",
    "\n",
    "#  e. Run inference \n",
    "with h5py.File(patch_features_path, 'r') as f:\n",
    "    coords = f['coords'][:]\n",
    "    patch_features = f['features'][:]\n",
    "    coords_attrs = dict(f['coords'].attrs)\n",
    "\n",
    "batch = {'features': torch.from_numpy(patch_features).float().to(device).unsqueeze(0)}\n",
    "logits, attention = model(batch, return_raw_attention=True)\n",
    "\n",
    "# f. generate heatmap\n",
    "heatmap_save_path = visualize_heatmap(\n",
    "    wsi=slide,\n",
    "    scores=attention.cpu().numpy().squeeze(),  \n",
    "    coords=coords,\n",
    "    vis_level=1,\n",
    "    patch_size_level0=coords_attrs['patch_size_level0'],\n",
    "    normalize=True,\n",
    "    num_top_patches_to_save=10,\n",
    "    output_dir=job_dir\n",
    ")\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "trident",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
